import networkx as nx
import pulp
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.patches import Circle
from matplotlib.font_manager import FontProperties
font = FontProperties(fname='resources/SimHei.ttf', size=12)  # 替换为实际字体路径

# ======================
# 1. 管网拓扑生成（随机数据）
# ======================
def generate_gas_network(num_valves=8, num_users=10, num_leaks=3):
    """生成随机燃气管网拓扑"""
    G = nx.DiGraph()
  
    # 添加节点类型：气源(0), 阀门(1), 用户(2), 泄漏点(3)
    nodes = []
    nodes.append((0, {'type': 'source', 'pressure': 0.8}))  # 气源节点
  
    for i in range(1, num_valves + 1):
        nodes.append((i, {'type': 'valve', 'operation_time': np.random.choice([60, 300], p=[0.3, 0.7])}))  # 操作时间(秒)
  
    for i in range(num_valves + 1, num_valves + num_users + 1):
        nodes.append((i, {'type': 'user', 'weight': np.random.choice([1, 5, 10], p=[0.7, 0.2, 0.1])}))  # 用户权重
  
    leak_nodes = np.random.choice(range(num_valves + num_users + 1, num_valves + num_users + num_leaks + 1), 
                                 num_leaks, replace=False)
    for i in leak_nodes:
        nodes.append((i, {'type': 'leak', 'leak_rate': np.random.uniform(0.5, 5.0)}))  # 泄漏率(m³/s)
  
    G.add_nodes_from(nodes)

    num_nodes = num_valves + num_users + num_leaks
  
    # 添加管道连接（边）
    for i in range(1, num_nodes):
        j = np.random.randint(0, num_nodes)
        if i != j and not G.has_edge(i, j):
            diameter = np.random.choice([150, 300, 500])  # mm
            length = np.random.randint(100, 1000)  # 米
            G.add_edge(i, j, diameter=diameter, length=length, pressure=np.random.uniform(0.2, 0.6))
            G.add_edge(j, i, diameter=diameter, length=length, pressure=np.random.uniform(0.2, 0.6))  # 双向管道
  
    # 确保气源连接到网络
    for i in range(1, min(5, num_nodes)):
        G.add_edge(0, i, diameter=500, length=200, pressure=0.8)
  
    return G

# 生成管网
np.random.seed(42)
G = generate_gas_network()

# ======================
# 2. 最小割算法（隔离泄漏点）
# ======================
def find_min_cut(G):
    """使用最小割算法确定需要关闭的阀门集合"""
    # 准备最小割计算
    source_node = 0  # 气源节点
    leak_nodes = [n for n, attr in G.nodes(data=True) if attr['type'] == 'leak']
  
    # 创建用于最小割计算的副本图
    H = nx.DiGraph()
    for u, v, attr in G.edges(data=True):
        # 阀门边的容量设为1（需要关闭），非阀门边容量设为无穷大
        cap = 1 if G.nodes[u].get('type') == 'valve' or G.nodes[v].get('type') == 'valve' else float('inf')
        H.add_edge(u, v, capacity=cap, **attr)
  
    # 计算最小割
    min_cut_value, partition = nx.minimum_cut(H, source_node, leak_nodes[0])
    reachable, non_reachable = partition
  
    # 确定需要关闭的阀门
    cut_edges = []
    for u, nbrs in ((n, H[n]) for n in reachable):
        for v, attr in nbrs.items():
            if v in non_reachable:
                if attr['capacity'] < float('inf'):
                    cut_edges.append((u, v))
    print (cut_edges)
    return set([u for u, v in cut_edges])

# 获取最小阀门集合
min_cut_valves = find_min_cut(G)
print (min_cut_valves)
#exit()
print(f"最小割算法确定的阀门集合: {min_cut_valves}")

# ======================
# 3. 多目标优化（PuLP实现）
# ======================
def optimize_valve_selection(G, candidate_valves):
    """使用线性规划优化阀门选择"""
    # 1. 问题初始化
    prob = pulp.LpProblem("Gas_Valve_Optimization", pulp.LpMinimize)
  
    # 2. 决策变量（是否关闭每个阀门）
    valve_vars = pulp.LpVariable.dicts("Valve", candidate_valves, cat='Binary')
  
    # 3. 目标函数参数（权重系数）
    alpha = 0.6  # 用户影响权重
    beta = 0.3   # 泄漏风险权重
    gamma = 0.1  # 操作时间权重
  
    # 4. 目标函数组件
    # a. 用户影响（停气用户总权重）
    user_impact = pulp.lpSum([
        G.nodes[node]['weight'] 
        for node in G.nodes 
        if G.nodes[node].get('type') == 'user'
        # 简化的连通性检查（实际应用需DFS/BFS）
        if any(valve in candidate_valves for valve in [node])
    ])
  
    # b. 泄漏风险（泄漏速率×预计持续时间）
    leak_risk = pulp.lpSum([
        G.nodes[node]['leak_rate'] * 300  # 假设平均持续300秒
        for node in G.nodes 
        if G.nodes[node].get('type') == 'leak'
    ])
  
    # c. 操作时间（阀门操作总时间）
    operation_time = pulp.lpSum([
        valve_vars[valve] * G.nodes[valve]['operation_time']
        for valve in candidate_valves
    ])
  
    # 5. 组合目标函数
    prob += alpha * user_impact + beta * leak_risk + gamma * operation_time
  
    # 6. 约束条件：必须隔离所有泄漏点（最小割已保证）
  
    # 7. 求解
    prob.solve(pulp.PULP_CBC_CMD(msg=False))
  
    # 8. 提取结果
    solution = [valve for valve in candidate_valves if pulp.value(valve_vars[valve]) == 1]
    print(f"优化后的阀门关闭方案: {solution}")
    print(f"预计操作时间: {pulp.value(operation_time)}秒")
    print(f"用户影响值: {pulp.value(user_impact)}")
    print(f"泄漏风险值: {pulp.value(leak_risk)}")
  
    return solution

# 执行优化
optimized_valves = optimize_valve_selection(G, min_cut_valves)

# ======================
# 4. 结果可视化
# ======================
def visualize_network(G, solution_valves):
    """可视化管网及关阀方案"""
    plt.figure(figsize=(14, 10))
  
    # 节点定位
    pos = nx.spring_layout(G, seed=42)
  
    # 按类型定义节点颜色
    node_colors = []
    node_sizes = []
    for node in G.nodes:
        attr = G.nodes[node]
        if attr['type'] == 'source':
            node_colors.append('lime')
            node_sizes.append(500)
        elif attr['type'] == 'valve':
            node_colors.append('lightblue' if node not in solution_valves else 'orange')
            node_sizes.append(400)
        elif attr['type'] == 'user':
            node_colors.append('gold')
            node_sizes.append(300)
        else:  # leak
            node_colors.append('red')
            node_sizes.append(450)
  
    # 绘制节点
    nx.draw_networkx_nodes(G, pos, node_size=node_sizes, node_color=node_colors, alpha=0.9)
  
    # 绘制边（突出显示需要关闭的阀门）
    edge_colors = []
    edge_widths = []
    for u, v in G.edges():
        if u in solution_valves or v in solution_valves:
            edge_colors.append('red')
            edge_widths.append(3.0)
        else:
            edge_colors.append('gray')
            edge_widths.append(1.5)
  
    nx.draw_networkx_edges(G, pos, edge_color=edge_colors, width=edge_widths, arrows=True, arrowstyle='-|>')
  
    # 节点标签
    labels = {}
    for node in G.nodes:
        attr = G.nodes[node]
        if attr['type'] == 'valve':
            labels[node] = f"V{node}\n({attr['operation_time']}s)"
        elif attr['type'] == 'user':
            labels[node] = f"U{node}\nW:{attr['weight']}"
        elif attr['type'] == 'leak':
            labels[node] = f"L{node}\n{attr['leak_rate']:.1f}m³/s"
        else:
            labels[node] = f"S{node}"
  
    nx.draw_networkx_labels(G, pos, labels, font_size=9)
  
    # 管道属性标注
    for (u, v), attr in G.edges.items():
        plt.text((pos[u][0]+pos[v][0])/2, 
                 (pos[u][1]+pos[v][1])/2, 
                 f"Ø{attr['diameter']}mm\n{attr['pressure']:.1f}MPa", 
                 fontsize=7, color='blue')
  
    # 图例
    legend_elements = [
        plt.Line2D([0], [0], marker='o', color='w', label='气源', markerfacecolor='lime', markersize=10),
        plt.Line2D([0], [0], marker='o', color='w', label='阀门(关闭)', markerfacecolor='orange', markersize=10),
        plt.Line2D([0], [0], marker='o', color='w', label='阀门(开启)', markerfacecolor='lightblue', markersize=10),
        plt.Line2D([0], [0], marker='o', color='w', label='用户', markerfacecolor='gold', markersize=10),
        plt.Line2D([0], [0], marker='o', color='w', label='泄漏点', markerfacecolor='red', markersize=10),
        plt.Line2D([0], [0], color='red', lw=2, label='需关闭阀门'),
        plt.Line2D([0], [0], color='gray', lw=1, label='普通管道')
    ]
  
    plt.legend(handles=legend_elements, loc='upper right', prop=font)
    plt.title("燃气管网泄漏关阀优化方案", fontsize=14, fontproperties=font)
    plt.tight_layout()
    plt.savefig('gas_network_optimization.png', dpi=300)
    plt.show()

# 可视化结果
visualize_network(G, optimized_valves)